// Prevents additional console window on Windows in release, DO NOT REMOVE!!
#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]

use std::sync::OnceLock;
use http::{header::*, status::StatusCode};

use regex::Regex;
use review_app::{Dataset, parse_dataset};
use serde::Serialize;
use tauri::http::{self, Response};
use url::Url;

static DATASET: OnceLock<Dataset> = OnceLock::new();

#[derive(Debug, thiserror::Error)]
enum Error {
  #[error(transparent)]
  Io(#[from] std::io::Error)
}

// we must manually implement serde::Serialize
impl serde::Serialize for Error {
  fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
  where
    S: serde::ser::Serializer,
  {
    serializer.serialize_str(self.to_string().as_ref())
  }
}

// Learn more about Tauri commands at https://tauri.app/v1/guides/features/command
#[tauri::command]
fn greet(name: &str) -> String {
    format!("Hello, {}! You've been greeted from Rust!", name)
}

#[tauri::command]
fn read_file(path: &str) -> Result<Vec<u8>, Error> {
    let output = std::fs::read(path)?;
    Ok(output)
}

#[tauri::command]
fn get_video_id() -> Vec<u32> {
    let dataset = DATASET.get().expect("get dataset failed");
    let ret = dataset.videos
        .keys()
        .copied()
        .collect::<Vec<_>>();
    ret
}

#[tauri::command]
fn video_frame_count(id: u32) -> usize {
    let dataset = DATASET.get().expect("get dataset failed");
    if let Some(v) = dataset.videos.get(&id) {
        v.len()
    } else {
        0
    }
}

#[tauri::command]
fn read_frame(video_id: u32, frame_id: u32) -> Result<Vec<u8>, Error> {
    let dataset = DATASET.get().expect("get dataset failed");
    if let Some(v) = dataset.videos.get(&video_id) {
        let frame_path = &v[(frame_id - 1) as usize];
        let ret = std::fs::read(frame_path)?;
        Ok(ret)
    } else {
        Ok(vec![])
    }
}

#[tauri::command]
fn load_label() -> Result<String, Error> {
    let dataset = DATASET.get().expect("get dataset failed");
    Ok(dataset.load_label())
}

#[tauri::command]
fn save_label(content: String) -> Result<(), Error> {
    let dataset = DATASET.get().expect("get dataset failed");
    dataset.save_label(content);
    Ok(())
}

#[derive(Serialize, Debug)]
struct Rot {
    raw: [f32; 4],
    pred: Option<[f32; 4]>
}

#[tauri::command]
fn frame_raw(video_id: u64, frame_id: u64) -> Rot {
    let dataset = DATASET.get().expect("get dataset failed");
    let raw = dataset.label.fetch_raw(video_id, frame_id - 1);
    let pred = dataset.infer.as_ref().and_then(|i| i.find_pred(video_id, frame_id - 1));
    Rot {
        raw,
        pred,
    }
}


fn response_frame(req: &http::Request) -> Result<Response, Box<dyn std::error::Error>> {
    let uri = req.uri();
    let url = Url::parse(uri)?;
    let path = url.path();
    // println!("request path: {}", path);

    let pattern_re = Regex::new(r"/(\d+)/(\d+)$").unwrap();
    let mut resp = if let Some(capture) = pattern_re.captures(path) {
        let video = capture.get(1).expect("can't find video id");
        let vid = u32::from_str_radix(&video.as_str(), 10).expect("invalid video id");

        let frame = capture.get(2).expect("can't find frame id");
        let fid = u32::from_str_radix(&frame.as_str(), 10).expect("invalid frame id");

        // println!("response frame with {}/{}", vid, fid);
        let dataset = DATASET.get().expect("get dataset failed");

        if let Some(v) = dataset.videos.get(&vid) {
            let frame_path = &v[(fid - 1) as usize];
            let content = std::fs::read(frame_path)?;
            let mut resp = Response::new(content);
            resp.set_status(StatusCode::OK);
            resp
        } else {
            let mut resp = Response::default();
            resp.set_status(StatusCode::NOT_FOUND);
            resp
        }
    } else {
        let mut resp = Response::default();
        resp.set_status(StatusCode::NOT_FOUND);
        resp
    };
    resp.headers_mut().insert(CONTENT_TYPE, HeaderValue::from_static("image/jpeg"));
    resp.headers_mut().insert(ACCESS_CONTROL_ALLOW_ORIGIN, HeaderValue::from_static("*"));
    println!("load a frame");
    Ok(resp)
}

fn main() {

    let args: Vec<String> = std::env::args().into_iter().collect();
    if args.len() < 2 {
        eprintln!("missing args, usage [ready dataset] [infer npz] or [ready dataset]");
        return;
    }

    let dataset_path = &args[1];
    let infer_path = if args.len() == 3 {
        Some(args[2].as_str())
    } else {
        None
    };

    // let infer_path = Some(INFER_PATH);
    let dataset = parse_dataset(dataset_path, infer_path).expect("parse dataset failed");

    DATASET.get_or_init(|| {
        dataset
    });

    tauri::Builder::default()
        .invoke_handler(tauri::generate_handler![
            greet,
            read_file,
            get_video_id,
            video_frame_count,
            read_frame,
            frame_raw,
            load_label,
            save_label,
        ])
        .register_uri_scheme_protocol("frame", move |_app, req| {
            response_frame(req)
        })
        .run(tauri::generate_context!())
        .expect("error while running tauri application");
}
